#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.0
# @author : cd
# @Email : 19688513@qq.com
# @Project : TradeDateManager
# @File : TradeDateManager.py
# @Software: PyCharm
# @time: 2025/6/11 15:38
# @description : 交易日管理器
"""

import logging
from typing import List, Tuple, Optional
from datetime import datetime, date
from diskcache import Cache
import pandas as pd
from bisect import bisect_left, bisect_right
import random
import time

from AKShareDataFetcher import AKShareDataFetcher

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


class TradeDateManager:
    def __init__(self, cache_directory: str = 'trade_dates_cache'):
        """
        初始化交易日管理器
        :param cache_directory: 缓存目录路径
        """
        self.cache = Cache(cache_directory)
        logging.info(f"交易日管理器初始化完成，缓存目录: {cache_directory}")
        # 初始化交易日列表
        self._trade_dates = self.read_trade_dates()
        # 确保日期排序
        if self._trade_dates:
            self._trade_dates.sort()

    def _format_date(self, input_date, date_format: str = '%Y%m%d') -> str:
        """
        格式化日期为指定格式的字符串
        :param input_date: 日期对象或字符串
        :param date_format: 返回日期的格式，默认为 'YYYYMMDD'
        :return: 格式化后的日期字符串
        """
        if input_date is None:
            raise ValueError("输入日期不能为None")

        if isinstance(input_date, date):
            return input_date.strftime(date_format)
        elif isinstance(input_date, str):
            # 尝试解析不同格式的日期字符串
            for fmt in ('%Y-%m-%d', '%Y%m%d', '%Y/%m/%d', '%Y%m%d', '%Y-%m-%d %H:%M:%S'):
                try:
                    parsed_date = datetime.strptime(input_date, fmt)
                    return parsed_date.strftime(date_format)
                except ValueError:
                    continue
            # 尝试直接返回前8位数字
            if input_date.strip().isdigit() and len(input_date.strip()) >= 8:
                return input_date.strip()[:8]
            raise ValueError(f"无法识别的日期格式: {input_date}")
        else:
            raise TypeError(f"不支持的日期类型: {type(input_date)}")

    def read_trade_dates(self) -> List[str]:
        """
        从缓存读取交易日期数据，并返回按时间排序的日期列表
        :return: 排序后的交易日期列表 (YYYYMMDD格式)
        """
        if 'trade_dates' in self.cache:
            dates = self.cache['trade_dates']
            # 确保格式统一为YYYYMMDD
            try:
                formatted_dates = [self._format_date(d) for d in dates]
                # 确保返回排序后的列表
                return sorted(formatted_dates)
            except Exception as e:
                logging.error(f"格式化缓存日期时出错: {e}")
                return []
        else:
            logging.warning("缓存中没有交易日期数据")
            return []

    def is_trade_day(self, check_date) -> bool:
        """
        判断给定日期是否为交易日
        :param check_date: 要检查的日期
        :return: 如果是交易日返回True，否则返回False
        """
        if check_date is None:
            raise ValueError("日期不能为空")

        try:
            formatted_date = self._format_date(check_date)
            return formatted_date in self._trade_dates
        except Exception as e:
            logging.error(f"检查交易日时出错: {e}")
            return False

    def get_nearest_trade_date(self, target_date, direction: str = 'previous') -> Optional[str]:
        """
        获取最近的有效交易日
        :param target_date: 目标日期
        :param direction: 查找方向 ('previous', 'next', 'closest')
        :return: 最近的交易日或None (YYYYMMDD格式)
        """
        try:
            formatted_date = self._format_date(target_date)
        except Exception as e:
            logging.error(f"格式化目标日期时出错: {e}")
            return None

        trade_dates = self._trade_dates

        if not trade_dates:
            logging.warning("交易日期列表为空")
            return None

        # 如果目标日期本身就是交易日
        if formatted_date in trade_dates:
            return formatted_date

        # 使用二分查找找到插入位置
        idx = bisect_left(trade_dates, formatted_date)

        if direction == 'previous':
            if idx > 0:
                return trade_dates[idx - 1]
            logging.warning(f"找不到 {formatted_date} 之前的交易日")
            return None
        elif direction == 'next':
            if idx < len(trade_dates):
                return trade_dates[idx]
            logging.warning(f"找不到 {formatted_date} 之后的交易日")
            return None
        elif direction == 'closest':
            if idx == 0:
                return trade_dates[0]
            if idx == len(trade_dates):
                return trade_dates[-1]

            # 计算前后交易日的距离
            prev_date = trade_dates[idx - 1]
            next_date = trade_dates[idx]
            try:
                prev_diff = (datetime.strptime(formatted_date, '%Y%m%d') -
                             datetime.strptime(prev_date, '%Y%m%d')).days
                next_diff = (datetime.strptime(next_date, '%Y%m%d') -
                             datetime.strptime(formatted_date, '%Y%m%d')).days

                return prev_date if prev_diff <= next_diff else next_date
            except Exception as e:
                logging.error(f"计算日期差时出错: {e}")
                return prev_date  # 默认返回前一个
        else:
            raise ValueError(f"无效的方向参数: {direction}。可选值: 'previous', 'next', 'closest'")

    def get_trade_dates_between(self, start_date, end_date) -> List[str]:
        """
        获取两个日期之间的所有交易日
        :param start_date: 开始日期
        :param end_date: 结束日期
        :return: 日期范围内的交易日列表 (YYYYMMDD格式)
        """
        try:
            start_fmt = self._format_date(start_date)
            end_fmt = self._format_date(end_date)
        except Exception as e:
            logging.error(f"格式化日期范围时出错: {e}")
            return []

        trade_dates = self._trade_dates

        if not trade_dates:
            logging.warning("交易日期列表为空")
            return []

        # 检查日期范围是否有效
        if start_fmt > end_fmt:
            logging.warning(f"开始日期 {start_fmt} 晚于结束日期 {end_fmt}")
            return []

        # 使用二分查找确定范围
        start_idx = bisect_left(trade_dates, start_fmt)
        end_idx = bisect_right(trade_dates, end_fmt) - 1

        if start_idx > end_idx:
            logging.info(f"在 {start_fmt} 和 {end_fmt} 之间没有找到交易日")
            return []

        return trade_dates[start_idx:end_idx + 1]

    def get_trade_window(self, target_date, days_before: int = 0, days_after: int = 0) -> Tuple[str, str]:
        """
        获取以目标日期为中心的交易窗口
        :param target_date: 目标日期
        :param days_before: 目标日期前的交易日数量
        :param days_after: 目标日期后的交易日数量
        :return: (开始日期, 结束日期) (YYYYMMDD格式)
        """
        try:
            formatted_date = self._format_date(target_date)
        except Exception as e:
            logging.error(f"格式化目标日期时出错: {e}")
            return None, None

        trade_dates = self._trade_dates

        if not trade_dates:
            logging.warning("交易日期列表为空")
            return None, None

        try:
            # 尝试找到目标日期的索引
            if formatted_date in trade_dates:
                idx = trade_dates.index(formatted_date)
            else:
                # 如果目标日期不是交易日，找到最近的交易日
                nearest_date = self.get_nearest_trade_date(formatted_date, 'closest')
                if not nearest_date:
                    logging.warning(f"找不到 {formatted_date} 附近的交易日")
                    return None, None
                idx = trade_dates.index(nearest_date)

            start_idx = max(0, idx - days_before)
            end_idx = min(len(trade_dates) - 1, idx + days_after)

            return trade_dates[start_idx], trade_dates[end_idx]
        except ValueError:
            logging.error(f"在交易日期列表中找不到日期")
            return None, None
        except Exception as e:
            logging.error(f"获取交易窗口时出错: {e}")
            return None, None

    def get_trade_dates(self, start_date: str, end_date: str) -> List[str]:
        """
        获取指定日期范围内的交易日
        :param start_date: 开始日期 (任意格式)
        :param end_date: 结束日期 (任意格式)
        :return: 交易日列表 (YYYYMMDD格式)
        """
        return self.get_trade_dates_between(start_date, end_date)

    def get_start_end_date(self, trade_date: str, y_day: int, date_format: str = '%Y%m%d') -> Tuple[str, str]:
        """
        根据给定的交易日期和天数，计算起始日期和结束日期
        :param trade_date: 交易日期 (字符串或日期对象)
        :param y_day: 要回溯的天数
        :param date_format: 返回日期的格式 (默认为YYYYMMDD)
        :return: (起始日期, 结束日期)
        """
        # 检查回溯天数
        if y_day < 0:
            logging.warning(f"回溯天数不能为负数: {y_day}，已自动设为0")
            y_day = 0

        try:
            # 格式化输入日期
            formatted_date = self._format_date(trade_date)
        except Exception as e:
            logging.error(f"格式化交易日期时出错: {e}")
            return None, None

        # 获取最近的交易日作为结束日期
        end_date = self.get_nearest_trade_date(formatted_date, 'previous')
        if not end_date:
            logging.error(f"找不到 {formatted_date} 附近的交易日")
            return None, None

        if not self._trade_dates:
            logging.warning("交易日期列表为空")
            return None, None

        try:
            # 找到结束日期的索引
            end_idx = self._trade_dates.index(end_date)
            # 计算起始日期的索引
            start_idx = max(0, end_idx - y_day)
            start_date = self._trade_dates[start_idx]

            # 返回格式化后的日期
            return self._format_date(start_date, date_format), self._format_date(end_date, date_format)
        except ValueError:
            logging.error(f"在交易日期列表中找不到 {end_date}")
            return None, None
        except Exception as e:
            logging.error(f"计算起始结束日期时出错: {e}")
            return None, None

    def insert_trade_dates(self, trade_dates: List[str]):
        """
        将交易日期插入缓存，忽略重复的日期并保持排序
        :param trade_dates: 要插入的交易日期列表 (任意格式)
        """
        if not trade_dates:
            logging.warning("尝试插入空的交易日期列表")
            return

        # 格式化所有日期为YYYYMMDD
        formatted_dates = []
        for d in trade_dates:
            try:
                formatted_dates.append(self._format_date(d))
            except Exception as e:
                logging.warning(f"跳过无效日期 {d}: {e}")

        if not formatted_dates:
            logging.warning("没有有效的日期可插入")
            return

        # 创建日期集合用于快速查找
        current_set = set(self._trade_dates)

        # 找出新日期
        new_dates = [d for d in formatted_dates if d not in current_set]

        if new_dates:
            # 合并并排序所有日期
            self._trade_dates = sorted(set(self._trade_dates) | set(formatted_dates))
            self.cache['trade_dates'] = self._trade_dates
            logging.info(f"成功插入 {len(new_dates)} 个新交易日期，总计 {len(self._trade_dates)} 个交易日。")

            # 确保内部状态更新
            self.load_trade_dates()
        else:
            logging.info("没有新的交易日期需要插入。")

    def get_cache_stats(self) -> dict:
        """获取缓存统计信息"""
        stats = self.cache.stats()
        return {
            'total_items': len(self.cache),
            'trade_dates_count': len(self._trade_dates),
            'cache_directory': self.cache.directory
        }

    def clear_cache(self):
        """清空缓存中的交易日期"""
        if 'trade_dates' in self.cache:
            del self.cache['trade_dates']
            self._trade_dates = []
            logging.info("交易日期缓存已清空。")
        else:
            logging.warning("缓存中没有可清空的交易日期。")

    def close(self):
        """关闭缓存连接"""
        self.cache.close()
        logging.info("交易日管理器缓存已关闭")

    def __del__(self):
        """析构函数，关闭缓存连接"""
        self.close()

    def initialize_trade_dates(self, force_refresh: bool = False):
        """
        获取并插入交易日期
        :param force_refresh: 是否强制刷新数据
        """
        # 如果缓存中已有数据且不需要强制刷新，则跳过
        if 'trade_dates' in self.cache and not force_refresh:
            logging.info("交易日期已存在缓存中，跳过初始化。")
            # 确保内部状态更新
            self.load_trade_dates()
            return

        logging.info("开始获取交易日期...")
        data_fetcher = AKShareDataFetcher(max_retries=3, retry_delay=5, request_delay=1)

        # 获取交易日期
        trade_dates = data_fetcher.fetch_trade_dates(max_retries=3, retry_delay=3, request_delay=0.5)

        if trade_dates is None or trade_dates.empty:
            logging.error("获取交易日期失败")
            return

        # 转换交易日期为字符串格式
        try:
            trade_dates['trade_date'] = pd.to_datetime(trade_dates['trade_date'], errors='coerce')
            valid_dates = trade_dates['trade_date'].dt.strftime('%Y%m%d').dropna().tolist()
        except Exception as e:
            logging.error(f"处理交易日期时出错: {e}")
            return

        if not valid_dates:
            logging.error("未找到有效的交易日期")
            return

        # 将有效的日期插入到交易日历中
        self.insert_trade_dates(valid_dates)
        logging.info(f"交易日期初始化完成，共添加 {len(valid_dates)} 个交易日。")

    def load_trade_dates(self):
        """
        确保交易日历加载到内存
        此方法是为了兼容性添加的
        """
        # 如果内部状态为空，重新加载
        if not self._trade_dates:
            self._trade_dates = self.read_trade_dates()
            if self._trade_dates:
                self._trade_dates.sort()
        logging.info(f"交易日历已加载到内存，共 {len(self._trade_dates)} 个交易日")

    def get_previous_trade_days(self, target_date: str, n_days: int) -> List[str]:
        """
        获取目标日期前n个交易日
        :param target_date: 目标日期 (YYYYMMDD)
        :param n_days: 要获取的交易日数量
        :return: 前n个交易日列表（最近的在最后）
        """
        nearest_date = self.get_nearest_trade_date(target_date, direction='previous')
        if not nearest_date:
            return []

        # 获取以nearest_date为中心，往前找n_days个交易日
        try:
            start_idx = max(0, self._trade_dates.index(nearest_date) - n_days)
            end_idx = self._trade_dates.index(nearest_date)
            return self._trade_dates[start_idx:end_idx]
        except ValueError:
            logging.error("找不到指定日期在交易日历中的索引")
            return []


# 示例用法
if __name__ == "__main__":
    # 创建交易日管理器
    trade_manager = TradeDateManager()

    # 加载交易日历
    trade_manager.load_trade_dates()

    # 初始化交易日期
    trade_manager.initialize_trade_dates()

    # 测试获取起始和结束日期
    test_date = "2025-06-08"  # 假设这是一个非交易日
    days_back = 2
    start_date, end_date = trade_manager.get_start_end_date(test_date, days_back)

    print(f"\n🔍 测试获取起始结束日期:")
    print(f"输入日期: {test_date}, 回溯天数: {days_back}")
    print(f"起始日期: {start_date}, 结束日期: {end_date}")

    # 验证结果
    if start_date and end_date:
        trade_dates = trade_manager.get_trade_dates_between(start_date, end_date)
        actual_days = len(trade_dates)
        print(f"实际包含天数: {actual_days} (预期: {days_back + 1})")

        # 打印日期范围内的交易日
        print(f"日期范围内的交易日: {trade_dates}")

    # 测试不同日期格式
    test_formats = [
        "2025-06-06",
        "20250606",
        date(2025, 6, 6),
        "2025/06/06",
        "2025-06-06 10:30:00",
        "invalid_date"
    ]

    print("\n🔄 日期格式兼容性测试:")
    for fmt in test_formats:
        try:
            result = trade_manager.get_nearest_trade_date(fmt, 'closest')
            print(f"输入格式: {fmt} -> 最近交易日: {result}")
        except Exception as e:
            print(f"输入格式: {fmt} -> 错误: {str(e)}")